import numpy as np
import gym
import time
from nes_py.wrappers import JoypadSpace
import gym_super_mario_bros
from gym.spaces import Box
from gym import Wrapper
from gym_super_mario_bros.actions import SIMPLE_MOVEMENT, COMPLEX_MOVEMENT, RIGHT_ONLY
from marioEnv import create_train_env
from marioEnv import ACTIONS
import tensorflow as tf
from tensorflow import keras
from policy_model import Policy

policy = Policy()
policy.load('policy.h5')

env = create_train_env(1, 1, True)
state = env.reset()

while True:
	env.render()
	#time.sleep(1./60.)
	# get input
	state.shape = (1,)+state.shape
	action, info = policy.sample_action(state)
	print(info)
	# step
	state, reward, done, info = env.step(action)
	if done:
		state = env.reset()
	time.sleep(1./60.)
	#input()